import pandas as pd
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
from matplotlib.colors import to_hex

# 假设df是已经加载的DataFrame
# 示例数据，实际应用中您会从'input.xlsx'中加载

plt.rcParams['font.sans-serif'] = ['SimHei']                        # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False                          # 用来正常显示负号

def percentage_to_hex_opacity(percentage):
    # 将百分比转换为十进制的透明度值
    decimal_opacity = int(round(percentage * 255 / 100))
    # 将十进制透明度值转换为十六进制
    hex_opacity = format(decimal_opacity, '02X')
    return hex_opacity

def insert_newline_every_three_chars(s):
    return '\n'.join(s[i:i+3] for i in range(0, len(s), 3))

def create(df,
    # 圆环布局的层级半径定义
    radii = {1: 10, 2: 20, 3: 30},
    # 设置不同层级的圆半径和字体大小
    node_sizes = {1: 10000, 2: 6000, 3: 4000},  # 根据层级设置节点大小
    font_sizes = {1: 18, 2: 16, 3: 14},  # 根据层级设置字体大小
    flat_ratio = 2,
    alpha2 = 60,  # 透明度百分比
    alpha3 = 20,   
    output_path='test.png'
    ):
    # df = pd.read_excel('input.xlsx')
    df = df.sort_values(['c1', 'c2', 'c3']).reset_index(drop=True)

    # 应用这个函数到c3列
    df['c3'] = df['c3'].apply(insert_newline_every_three_chars)

    # 创建图
    G = nx.DiGraph()

    # 添加节点和边
    for _, row in df.iterrows():
        G.add_node(row['c1'], level=1, content=row['c1'])
        G.add_node(row['c2'], level=2, parent=row['c1'])
        G.add_node(row['c3'], level=3, parent=row['c2'])
        G.add_edge(row['c1'], row['c2'])
        G.add_edge(row['c2'], row['c3'])


    # 分配颜色给c1的不同内容
    base_colors = plt.get_cmap('tab10')
    color_map = {value: to_hex(base_colors(i % 10)) for i, value in enumerate(pd.unique(df['c1']))}

    # 为每个层级的节点分配位置
    # 开始时，先均匀分布c3层级的节点
    c3_nodes = [node for node in G.nodes() if G.nodes[node]['level'] == 3]
    
    angle_step = 2 * np.pi / len(c3_nodes)
    pos = {}
    for i, node in enumerate(c3_nodes):
        angle = i * angle_step
        pos[node] = (radii[3] * np.cos(angle) * flat_ratio, radii[3] * np.sin(angle))

    # 基于c3节点的位置计算c2节点的位置
    for node in set(df['c2']):
        children = [n for n in G.successors(node)]
        angle_list = [np.arctan2(pos[child][1], pos[child][0]) for child in children]
        max_angle = max(angle_list)
        min_angle = min(angle_list)
        if max_angle*min_angle<0:
            angle_list = [abs(x) for x in angle_list]
        avg_angle = np.mean(angle_list)
        pos[node] = (radii[2] * np.cos(avg_angle) * flat_ratio, radii[2] * np.sin(avg_angle))  # 横坐标乘以1.3

    # 基于c2节点的位置计算c1节点的位置
    for node in set(df['c1']):
        children = [n for n in G.successors(node)]
        avg_angle = np.mean([np.arctan2(pos[child][1], pos[child][0]) for child in children])
        pos[node] = (radii[1] * np.cos(avg_angle) * flat_ratio, radii[1] * np.sin(avg_angle))  # 横坐标乘以1.3

    # 设置节点颜色和透明度
    node_colors = []
    for node in G.nodes():
        level = G.nodes[node]['level']
        if level == 1:
            node_colors.append(color_map[G.nodes[node]['content']])
        elif level == 2:
            parent_content = G.nodes[G.nodes[node]['parent']]['content']
            node_colors.append(color_map[parent_content] + percentage_to_hex_opacity(alpha2))  # 0.8透明度
        else:  # level 3
            parent_node = G.nodes[node]['parent']
            grandparent_content = G.nodes[G.nodes[parent_node]['parent']]['content']
            node_colors.append(color_map[grandparent_content]  + percentage_to_hex_opacity(alpha3))  # 0.5透明度


    # 绘制图形
    fig_height = 8
    plt.figure(figsize=(flat_ratio*fig_height, fig_height))
    nx.draw(G, pos, with_labels=True, arrows=True, node_size=3000, node_color=node_colors, font_size=14, font_weight="bold", font_family='SimHei')





    plt.figure(figsize=(flat_ratio*fig_height, fig_height))

    # 绘制节点和边
    for level in node_sizes:
        # 筛选当前层级的节点
        nodes_at_level = [node for node in G.nodes() if G.nodes[node]['level'] == level]
        # 获取当前层级节点的颜色
        colors_at_level = [node_colors[i] for i, node in enumerate(G.nodes()) if G.nodes[node]['level'] == level]
        # 绘制当前层级的节点
        nx.draw_networkx_nodes(G, pos, nodelist=nodes_at_level, node_size=node_sizes[level], node_color=colors_at_level)

    # 绘制边，并将颜色设置为灰色
    nx.draw_networkx_edges(G, pos, edge_color="grey")

    # 单独绘制每个层级的标签以设置不同的字体大小
    for level in font_sizes:
        nodes_at_level = [node for node in G.nodes() if G.nodes[node]['level'] == level]
        labels_at_level = {node: node for node in nodes_at_level}
        nx.draw_networkx_labels(G, pos, labels=labels_at_level, font_size=font_sizes[level])


    # 在绘制完图形后，关闭当前轴的边框
    plt.gca().set_axis_off()

    # 调整子图参数，确保图像周围没有额外的边距
    plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)

    # 设置绘图边界
    plt.margins(0,0)

    plt.axis('equal')
    plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
    # plt.show()

if __name__=='__main__':
    create()


